#include "auto_navigation/mapping/rtk_ndt_mapper.hpp"

namespace auto_navigation {

RTKNDTMapper::RTKNDTMapper(ros::NodeHandle& nh, ros::NodeHandle& private_nh)
    : nh_(nh)
    , private_nh_(private_nh)
    , tf_listener_(tf_buffer_)
    , map_cloud_(new pcl::PointCloud<pcl::PointXYZ>)
    , current_cloud_(new pcl::PointCloud<pcl::PointXYZ>)
    , is_first_frame_(true)
    , has_rtk_data_(false)
{
    initializeParams();
    
    // 初始化订阅器（两种模式：同步和非同步）
    bool use_sync_mode = private_nh_.param<bool>("use_sync_mode", true);
    
    if (use_sync_mode) {
        // 使用同步模式
        sync_points_sub_.reset(new message_filters::Subscriber<sensor_msgs::PointCloud2>(nh_, "velodyne_points", 10));
        sync_rtk_odom_sub_.reset(new message_filters::Subscriber<nav_msgs::Odometry>(nh_, "rtk_odometry_node/odom", 10));
        
        sync_.reset(new message_filters::Synchronizer<SyncPolicy>(
            SyncPolicy(10), *sync_points_sub_, *sync_rtk_odom_sub_));
        sync_->registerCallback(boost::bind(&RTKNDTMapper::syncCallback, this, _1, _2));
        
        ROS_INFO("RTK-NDT Mapper initialized in synchronized mode");
    } else {
        // 使用非同步模式
        points_sub_ = nh_.subscribe("velodyne_points", 1, &RTKNDTMapper::pointCloudCallback, this);
        rtk_odom_sub_ = nh_.subscribe("rtk_odometry_node/odom", 1, &RTKNDTMapper::rtkOdomCallback, this);
        
        ROS_INFO("RTK-NDT Mapper initialized in non-synchronized mode");
    }
    
    // 初始化发布器
    map_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("rtk_ndt/map", 1);
    current_pose_pub_ = nh_.advertise<geometry_msgs::PoseStamped>("rtk_ndt/current_pose", 1);
    current_scan_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("rtk_ndt/current_scan", 1);
    
    // 初始化服务
    save_map_service_ = nh_.advertiseService("save_map", &RTKNDTMapper::saveMapCallback, this);
    
    // 初始化NDT参数
    ndt_.setTransformationEpsilon(ndt_epsilon_);
    ndt_.setStepSize(ndt_step_size_);
    ndt_.setResolution(ndt_resolution_);
    ndt_.setMaximumIterations(max_iterations_);
    
    // 初始化位姿
    current_pose_ = Eigen::Matrix4f::Identity();
    previous_pose_ = Eigen::Matrix4f::Identity();
    rtk_pose_ = Eigen::Matrix4f::Identity();
    
    ROS_INFO("RTK-NDT Mapper initialized");
}

RTKNDTMapper::~RTKNDTMapper() {
    std::string map_file = ros::package::getPath("auto_navigation") + "/maps/rtk_ndt_map.pcd";
    if (saveMap(map_file)) {
        ROS_INFO("Map saved automatically on shutdown to %s", map_file.c_str());
    }
}

void RTKNDTMapper::initializeParams() {
    // 从参数服务器读取参数，如果没有则使用默认值
    private_nh_.param("voxel_leaf_size", voxel_leaf_size_, 0.3);
    private_nh_.param("ndt_resolution", ndt_resolution_, 1.0);
    private_nh_.param("ndt_step_size", ndt_step_size_, 0.1);
    private_nh_.param("ndt_epsilon", ndt_epsilon_, 0.01);
    private_nh_.param("max_iterations", max_iterations_, 30);
    
    // RTK相关参数
    private_nh_.param("rtk_weight", rtk_weight_, 0.5);
    private_nh_.param("rtk_max_distance", rtk_max_distance_, 5.0);
    
    // 坐标系参数
    private_nh_.param<std::string>("map_frame_id", map_frame_id_, "map");
    private_nh_.param<std::string>("base_frame_id", base_frame_id_, "base_link");
    private_nh_.param<std::string>("lidar_frame_id", lidar_frame_id_, "velodyne");
    
    ROS_INFO("RTK-NDT Mapper parameters:");
    ROS_INFO("  voxel_leaf_size: %.2f", voxel_leaf_size_);
    ROS_INFO("  ndt_resolution: %.2f", ndt_resolution_);
    ROS_INFO("  ndt_step_size: %.2f", ndt_step_size_);
    ROS_INFO("  ndt_epsilon: %.4f", ndt_epsilon_);
    ROS_INFO("  max_iterations: %d", max_iterations_);
    ROS_INFO("  rtk_weight: %.2f", rtk_weight_);
    ROS_INFO("  rtk_max_distance: %.2f", rtk_max_distance_);
    ROS_INFO("  map_frame_id: %s", map_frame_id_.c_str());
    ROS_INFO("  base_frame_id: %s", base_frame_id_.c_str());
    ROS_INFO("  lidar_frame_id: %s", lidar_frame_id_.c_str());
}

void RTKNDTMapper::pointCloudCallback(const sensor_msgs::PointCloud2::ConstPtr& input_cloud) {
    // 非同步模式下的点云回调
    
    // 将ROS消息转换为PCL点云
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::fromROSMsg(*input_cloud, *cloud);
    
    // 降采样点云
    pcl::PointCloud<pcl::PointXYZ>::Ptr filtered_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    downsamplePointCloud(cloud, filtered_cloud, voxel_leaf_size_);
    
    // 保存当前点云
    *current_cloud_ = *filtered_cloud;
    
    // 如果是第一帧，初始化地图
    if (is_first_frame_) {
        *map_cloud_ = *filtered_cloud;
        is_first_frame_ = false;
        
        // 发布地图和当前点云
        publishMap();
        publishCurrentScan();
        
        ROS_INFO("First point cloud received, initializing map");
        return;
    }
    
    // 点云配准
    if (alignPointCloud()) {
        // 更新地图
        updateMap();
        
        // 发布地图、位姿和当前点云
        publishMap();
        publishPose();
        publishCurrentScan();
    }
}

void RTKNDTMapper::rtkOdomCallback(const nav_msgs::Odometry::ConstPtr& rtk_odom) {
    // 非同步模式下的RTK回调
    rtk_pose_ = odomToEigenMatrix(rtk_odom);
    has_rtk_data_ = true;
}

void RTKNDTMapper::syncCallback(const sensor_msgs::PointCloud2::ConstPtr& input_cloud, 
                               const nav_msgs::Odometry::ConstPtr& rtk_odom) {
    // 同步模式下的回调
    
    // 更新RTK位姿
    rtk_pose_ = odomToEigenMatrix(rtk_odom);
    has_rtk_data_ = true;
    
    // 处理点云
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::fromROSMsg(*input_cloud, *cloud);
    
    // 降采样点云
    pcl::PointCloud<pcl::PointXYZ>::Ptr filtered_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    downsamplePointCloud(cloud, filtered_cloud, voxel_leaf_size_);
    
    // 保存当前点云
    *current_cloud_ = *filtered_cloud;
    
    // 如果是第一帧，初始化地图
    if (is_first_frame_) {
        *map_cloud_ = *filtered_cloud;
        is_first_frame_ = false;
        
        // 使用RTK数据初始化位姿
        current_pose_ = rtk_pose_;
        previous_pose_ = current_pose_;
        
        // 发布地图和当前点云
        publishMap();
        publishCurrentScan();
        publishPose();
        
        ROS_INFO("First point cloud received, initializing map with RTK data");
        return;
    }
    
    // 点云配准
    if (alignPointCloud()) {
        // 使用RTK数据更新位姿
        updatePoseWithRTK();
        
        // 更新地图
        updateMap();
        
        // 发布地图、位姿和当前点云
        publishMap();
        publishPose();
        publishCurrentScan();
    }
}

bool RTKNDTMapper::alignPointCloud() {
    // 获取初始猜测
    Eigen::Matrix4f initial_guess;
    if (!getInitialGuess(initial_guess)) {
        ROS_WARN("Failed to get initial guess for NDT alignment");
        return false;
    }
    
    // 设置NDT输入
    ndt_.setInputSource(current_cloud_);
    ndt_.setInputTarget(map_cloud_);
    
    // 执行配准
    pcl::PointCloud<pcl::PointXYZ>::Ptr output_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    ndt_.align(*output_cloud, initial_guess);
    
    // 检查配准是否成功
    if (!ndt_.hasConverged()) {
        ROS_WARN("NDT alignment did not converge");
        return false;
    }
    
    // 获取变换矩阵
    Eigen::Matrix4f ndt_transform = ndt_.getFinalTransformation();
    
    // 更新位姿
    previous_pose_ = current_pose_;
    current_pose_ = ndt_transform;
    
    // 输出配准信息
    double fitness_score = ndt_.getFitnessScore();
    int iterations = ndt_.getFinalNumIteration();
    ROS_INFO("NDT alignment: score=%.4f, iterations=%d", fitness_score, iterations);
    
    return true;
}

void RTKNDTMapper::updatePoseWithRTK() {
    if (!has_rtk_data_) {
        ROS_WARN("No RTK data available for pose correction");
        return;
    }
    
    // 计算NDT位姿和RTK位姿之间的距离
    Eigen::Vector3f ndt_position(current_pose_(0, 3), current_pose_(1, 3), current_pose_(2, 3));
    Eigen::Vector3f rtk_position(rtk_pose_(0, 3), rtk_pose_(1, 3), rtk_pose_(2, 3));
    
    double distance = (ndt_position - rtk_position).norm();
    
    // 如果距离超过阈值，不使用RTK数据
    if (distance > rtk_max_distance_) {
        ROS_WARN("Distance between NDT and RTK poses (%.2f m) exceeds threshold (%.2f m), not using RTK data",
                distance, rtk_max_distance_);
        return;
    }
    
    // 使用加权平均融合NDT和RTK位姿
    Eigen::Matrix4f fused_pose = Eigen::Matrix4f::Identity();
    
    // 位置融合（加权平均）
    fused_pose(0, 3) = (1 - rtk_weight_) * current_pose_(0, 3) + rtk_weight_ * rtk_pose_(0, 3);
    fused_pose(1, 3) = (1 - rtk_weight_) * current_pose_(1, 3) + rtk_weight_ * rtk_pose_(1, 3);
    fused_pose(2, 3) = (1 - rtk_weight_) * current_pose_(2, 3) + rtk_weight_ * rtk_pose_(2, 3);
    
    // 旋转融合（简单方法：使用NDT的旋转）
    // 注意：更复杂的方法可以使用四元数插值
    fused_pose.block<3, 3>(0, 0) = current_pose_.block<3, 3>(0, 0);
    
    // 更新位姿
    current_pose_ = fused_pose;
    
    ROS_INFO("Pose corrected with RTK data (weight=%.2f, distance=%.2f m)", rtk_weight_, distance);
}

void RTKNDTMapper::updateMap() {
    // 将当前点云转换到地图坐标系
    pcl::PointCloud<pcl::PointXYZ>::Ptr transformed_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::transformPointCloud(*current_cloud_, *transformed_cloud, current_pose_);
    
    // 将转换后的点云添加到地图中
    *map_cloud_ += *transformed_cloud;
    
    // 对地图进行降采样以减少点云大小
    pcl::PointCloud<pcl::PointXYZ>::Ptr filtered_map(new pcl::PointCloud<pcl::PointXYZ>);
    downsamplePointCloud(map_cloud_, filtered_map, voxel_leaf_size_);
    
    // 更新地图
    *map_cloud_ = *filtered_map;
    
    ROS_INFO("Map updated, current size: %lu points", map_cloud_->size());
}

void RTKNDTMapper::publishMap() {
    // 发布地图点云
    sensor_msgs::PointCloud2 map_msg;
    pcl::toROSMsg(*map_cloud_, map_msg);
    map_msg.header.stamp = ros::Time::now();
    map_msg.header.frame_id = map_frame_id_;
    map_pub_.publish(map_msg);
}

void RTKNDTMapper::publishPose() {
    // 发布当前位姿
    geometry_msgs::PoseStamped pose_msg;
    pose_msg.header.stamp = ros::Time::now();
    pose_msg.header.frame_id = map_frame_id_;
    
    // 提取位置
    pose_msg.pose.position.x = current_pose_(0, 3);
    pose_msg.pose.position.y = current_pose_(1, 3);
    pose_msg.pose.position.z = current_pose_(2, 3);
    
    // 提取旋转（从旋转矩阵转换为四元数）
    Eigen::Matrix3f rotation = current_pose_.block<3, 3>(0, 0);
    Eigen::Quaternionf q(rotation);
    pose_msg.pose.orientation.x = q.x();
    pose_msg.pose.orientation.y = q.y();
    pose_msg.pose.orientation.z = q.z();
    pose_msg.pose.orientation.w = q.w();
    
    current_pose_pub_.publish(pose_msg);
    
    // 发布TF变换
    geometry_msgs::TransformStamped transform_stamped;
    transform_stamped.header.stamp = ros::Time::now();
    transform_stamped.header.frame_id = map_frame_id_;
    transform_stamped.child_frame_id = base_frame_id_;
    
    transform_stamped.transform.translation.x = current_pose_(0, 3);
    transform_stamped.transform.translation.y = current_pose_(1, 3);
    transform_stamped.transform.translation.z = current_pose_(2, 3);
    
    transform_stamped.transform.rotation.x = q.x();
    transform_stamped.transform.rotation.y = q.y();
    transform_stamped.transform.rotation.z = q.z();
    transform_stamped.transform.rotation.w = q.w();
    
    tf_broadcaster_.sendTransform(transform_stamped);
}

void RTKNDTMapper::publishCurrentScan() {
    // 发布当前点云
    sensor_msgs::PointCloud2 scan_msg;
    pcl::toROSMsg(*current_cloud_, scan_msg);
    scan_msg.header.stamp = ros::Time::now();
    scan_msg.header.frame_id = lidar_frame_id_;
    current_scan_pub_.publish(scan_msg);
}

bool RTKNDTMapper::getInitialGuess(Eigen::Matrix4f& initial_guess) {
    // 使用上一帧到当前帧的变换作为初始猜测
    if (!is_first_frame_) {
        // 计算上一帧到当前帧的变换
        Eigen::Matrix4f last_transform = previous_pose_.inverse() * current_pose_;
        
        // 使用上一帧的变换作为初始猜测
        initial_guess = current_pose_ * last_transform;
        
        return true;
    }
    
    // 如果是第一帧，使用单位矩阵作为初始猜测
    initial_guess = Eigen::Matrix4f::Identity();
    return true;
}

Eigen::Matrix4f RTKNDTMapper::odomToEigenMatrix(const nav_msgs::Odometry::ConstPtr& odom_msg) {
    // 将Odometry消息转换为Eigen矩阵
    Eigen::Matrix4f matrix = Eigen::Matrix4f::Identity();
    
    // 提取位置
    matrix(0, 3) = odom_msg->pose.pose.position.x;
    matrix(1, 3) = odom_msg->pose.pose.position.y;
    matrix(2, 3) = odom_msg->pose.pose.position.z;
    
    // 提取旋转（从四元数转换为旋转矩阵）
    Eigen::Quaternionf q(
        odom_msg->pose.pose.orientation.w,
        odom_msg->pose.pose.orientation.x,
        odom_msg->pose.pose.orientation.y,
        odom_msg->pose.pose.orientation.z
    );
    
    matrix.block<3, 3>(0, 0) = q.toRotationMatrix();
    
    return matrix;
}

void RTKNDTMapper::downsamplePointCloud(const pcl::PointCloud<pcl::PointXYZ>::Ptr& input_cloud,
                                      pcl::PointCloud<pcl::PointXYZ>::Ptr& output_cloud,
                                      double leaf_size) {
    // 使用体素网格滤波器降采样点云
    pcl::VoxelGrid<pcl::PointXYZ> voxel_filter;
    voxel_filter.setInputCloud(input_cloud);
    voxel_filter.setLeafSize(leaf_size, leaf_size, leaf_size);
    voxel_filter.filter(*output_cloud);
}

bool RTKNDTMapper::saveMapCallback(std_srvs::Trigger::Request& req, std_srvs::Trigger::Response& res) {
    // 保存地图服务回调
    std::string map_file = ros::package::getPath("auto_navigation") + "/maps/rtk_ndt_map_" + 
                          std::to_string(ros::Time::now().toSec()) + ".pcd";
    
    bool success = saveMap(map_file);
    
    res.success = success;
    if (success) {
        res.message = "Map saved successfully to " + map_file;
        ROS_INFO("%s", res.message.c_str());
    } else {
        res.message = "Failed to save map to " + map_file;
        ROS_ERROR("%s", res.message.c_str());
    }
    
    return true;
}

bool RTKNDTMapper::saveMap(const std::string& filename) {
    // 保存地图到PCD文件
    if (map_cloud_->empty()) {
        ROS_WARN("Map is empty, not saving");
        return false;
    }
    
    try {
        pcl::io::savePCDFileBinary(filename, *map_cloud_);
        ROS_INFO("Map saved to %s (%lu points)", filename.c_str(), map_cloud_->size());
        return true;
    } catch (const std::exception& e) {
        ROS_ERROR("Failed to save map: %s", e.what());
        return false;
    }
}

} // namespace auto_navigation 